package edu.cmu.lti.lm.dt;

import java.io.FileWriter;
import java.io.IOException;
import java.io.OutputStreamWriter;
import java.io.PrintWriter;
import java.util.Iterator;
import java.util.LinkedList;
import java.util.List;
import java.util.Set;

public class Trainer {
	private Set<Feature> features;
	private EntryCollection entries;
	public Trainer(Set<Feature> features, EntryCollection entries) {
		super();
		this.features = features;
		this.entries = entries;
	}
	public EntryCollection getEntries() {
		return entries;
	}
	public void setEntries(EntryCollection entries) {
		this.entries = entries;
	}
	
	
	protected TreeNode doTrain(EntryCollection ent, int level, String left) throws IOException{
		List<Pair<Question, Double>> l = new LinkedList<Pair<Question, Double>>();
		Iterator<Feature> it = features.iterator();
		double score = ent.getConditionalEntropy();
		Question res = null;
		double bef = score;
		PrintWriter wr = new PrintWriter(new OutputStreamWriter(new NullOutputStream()));
		//PrintWriter wr = new PrintWriter(new FileWriter("/people/qing0/G"+level+left));
		while(it.hasNext()){
			Feature ft = it.next();
			Pair<Question, Double> obtained = ft.getBestQuestion(ent,wr,bef);
			l.add(obtained);
			if(obtained.getSecond() < score){
				score = obtained.getSecond();
				res = obtained.getFirst();
			}
		}
		wr.close();
		if(res != null){
			Pair<EntryCollection,EntryCollection> r = 
				ent.divide(res);
			if(level>0){
				System.err.println("TRY TO RUN FURTHER ITERATIONS " + level);
				TreeNode nPos = doTrain(r.getFirst(),level - 1,left+"P");
				TreeNode nNeg = doTrain(r.getSecond(),level - 1,left+"N");
				return new TreeNode(nPos, nNeg, res);
			}else{
				// Train Model
				TreeNode w = new  TreeNode(null, null, res);
				w.posDist = EntryCollection.wordDistribute(r.getFirst());
				w.negDist = EntryCollection.wordDistribute(r.getSecond());
				return w;
			}
			//TODO: Stop criteria
		}
		return null;
	}
	
	public TreeNode train(int level) throws IOException{
		return doTrain(entries,level,"");
		// cross validation?!

		
	}
	
}
